package main

import (
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"math/rand"
	"net/http"
	"net/http/httputil"
	"net/url"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/gin-contrib/static"
	"github.com/gin-gonic/gin"
	"github.com/inu1255/brook"
	_ "github.com/jinzhu/gorm/dialects/sqlite"
	"github.com/mitchellh/mapstructure"
	"github.com/sirupsen/logrus"
	"github.com/urfave/cli"
)

type Config struct {
	Key   string `json:"key"`
	Token string `json:"token"`
	Port  int    `json:"port"`
}

type Server interface {
	Shutdown() error
	ListenAndServe() error
}

var logger = logrus.New()
var servers = make(map[int]Server)  // 所有的 brook server
var users = make([]*User, 0, 100)   // 用户列表
var port2user = make(map[int]*User) // port -> user
var port2user_lock = sync.Mutex{}
var changes = make([]*User, 0, 100)
var config = new(Config)
var lock = sync.RWMutex{}
var debug bool
var refreshing bool
var logfile string

func main() {
	logger.SetFormatter(&logrus.TextFormatter{
		TimestampFormat: "2006-01-02 03:04:05",
		FullTimestamp:   true,
	})
	app := cli.NewApp()
	app.Name = "Brook Panel"
	app.Version = "20200102"
	app.Usage = "Brook Web Panel"
	app.Flags = []cli.Flag{
		&cli.BoolFlag{
			Name:        "debug",
			Aliases:     []string{"d"},
			Usage:       "Enable debug",
			Destination: &debug,
		},
		&cli.StringFlag{
			Name:        "log",
			Aliases:     []string{"l"},
			Usage:       "log file",
			Destination: &logfile,
		},
	}
	app.Commands = []*cli.Command{
		&cli.Command{
			Name:  "start",
			Usage: "启动服务",
			Flags: []cli.Flag{
				&cli.StringFlag{
					Name:    "key",
					Aliases: []string{"k"},
					Usage:   "server key",
				},
				&cli.StringFlag{
					Name:    "port",
					Aliases: []string{"p"},
					Usage:   "server port",
					Value:   "9182",
				},
			},
			Action: func(c *cli.Context) error {
				if logfile != "" {
					f, err := os.OpenFile(logfile, os.O_CREATE|os.O_WRONLY, 0644)
					if err == nil {
						logger.Out = f
					}
				}
				gin.SetMode(gin.ReleaseMode)
				config.Port = c.Int("port")
				ip := getIP(c.String("key"))
				initDB()

				brook.Traffic.OnData(func(port, up, down int) {
					// 允许统计误差，以保证性能
					user := port2user[port]
					if user != nil {
						user.U += int64(up)
						user.D += int64(down)
						if user.StartAt > 0 && user.U+user.D > user.Limit {
							startUser(user)
						}
						now := time.Now()
						if now.Sub(time.Unix(user.UpdateAt/1000, 0)) > 5*time.Second {
							user.UpdateAt = now.UnixNano() / 1e6
							db.Model(user).Update(user)
						}
					} else {
						_, _ = httpPost("/node/traffic", gin.H{
							"port": port,
							"u":    up,
							"d":    down,
						})
					}
				})

				// 本地用户
				var rows []User
				db.Find(&rows)
				for i := range rows {
					users = append(users, &rows[i])
				}

				for _, user := range users {
					startUser(user)
				}

				t := time.NewTicker(12 * time.Hour)
				go func() {
					for {
						reloadUsers(3)
						<-t.C
					}
				}()

				r := gin.New()
				r.Use(gin.Recovery())
				if debug {
					r.Use(gin.Logger())
				}
				r.Use(static.Serve("/", static.LocalFile("public", false)))
				a := r.Group("/api")
				a.GET("/ports", func(c *gin.Context) {
					if c.Query("kento") == "brookird@fuck.qq.com" {
						var ports = make([]int, len(users))
						for _, user := range users {
							ports = append(ports, user.Port)
						}
						SendOk(c, gin.H{
							"ports": ports,
						})
					}
				})
				a.GET("/reload", func(c *gin.Context) {
					if refreshing {
						return
					}
					refreshing = true
					reloadUsers(5)
					refreshing = false
				})
				a.GET("/brook/getmyinfo", func(c *gin.Context) {
					id, err := strconv.Atoi(c.Query("id"))
					if err != nil {
						SendError(c, 400, "参数不正确")
						return
					}

					var row *User
					for _, item := range users {
						if item.ID == id {
							row = item
							break
						}
					}

					if row == nil {
						SendError(c, 404, "参数不正确")
						return
					}

					if c.Query("token") != row.Token {
						SendError(c, 404, "参数不正确")
						return
					}
					SendOk(c, row)
				})
				a.POST("/brook/setmyinfo", func(c *gin.Context) {
					var body = new(BrookBody)
					if err := c.ShouldBind(body); err != nil {
						SendError(c, 400, "参数错误: "+err.Error())
						return
					}

					var row *User
					for _, item := range users {
						if item.ID == body.ID {
							row = item
							break
						}
					}

					if row == nil {
						SendError(c, 404, "参数不正确")
						return
					}

					if body.Token != row.Token {
						SendError(c, 404, "参数不正确")
						return
					}

					if body.Passwd != "" || body.Port > 10000 && body.Port < 65535 {
						if body.Passwd != "" {
							row.Passwd = body.Passwd
						}
						if body.Port > 10000 && body.Port < 65535 {
							row.Port = body.Port
						}
						startUser(row)
					}

					SendOk(c, row)
				})
				a.GET("/users/list", func(c *gin.Context) {
					if c.Query("token") != config.Token {
						SendError(c, 403, "口令错误")
						return
					}
					SendOk(c, gin.H{
						"list":  users,
						"total": len(users),
					})
				})
				a.POST("/users/add", func(c *gin.Context) {
					if c.Query("token") != config.Token {
						SendError(c, 403, "口令错误")
						return
					}
					var user = new(User)
					if err := c.ShouldBind(user); err != nil {
						SendError(c, 400, "参数错误: "+err.Error())
						return
					}
					if servers[user.Port] != nil {
						SendError(c, 405, "端口被占用")
						return
					}
					if user.Passwd == "" {
						SendError(c, 400, "密码不能为空")
						return
					}
					user.CreateAt = time.Now().UnixNano() / 1e6
					startUser(user)
					db.Create(user)
					users = append(users, user)
					SendOk(c, user)
				})
				a.POST("/users/set", func(c *gin.Context) {
					if c.Query("token") != config.Token {
						SendError(c, 403, "口令错误")
						return
					}
					var user = new(User)
					if err := c.ShouldBind(user); err != nil {
						SendError(c, 400, "参数错误: "+err.Error())
						return
					}
					var row *User

					for _, item := range users {
						if item.ID == user.ID {
							row = item
							break
						}
					}

					if row == nil {
						SendError(c, 404, "用户不存在")
						return
					}
					if user.Port != row.Port && servers[user.Port] != nil {
						SendError(c, 405, "端口被占用")
						return
					}
					if user.Passwd == "" {
						SendError(c, 400, "密码不能为空")
						return
					}
					oldport := row.Port
					row.D = user.D
					row.U = user.U
					row.Limit = user.Limit
					row.ExpireAt = user.ExpireAt
					row.Passwd = user.Passwd
					row.Port = user.Port
					row.Speed = user.Speed
					row.Token = user.Token
					row.Enable = user.Enable
					row.Mark = user.Mark
					row.Domain = user.Domain
					startUser(row)
					db.Model(row).Update(row)
					if user.Port != oldport && user.Port > 0 {
						stopBrook(oldport)
					}
					SendOk(c, row)
				})
				a.GET("/users/del", func(c *gin.Context) {
					if c.Query("token") != config.Token {
						SendError(c, 403, "口令错误")
						return
					}
					id, err := strconv.Atoi(c.Query("id"))
					if err != nil {
						SendError(c, 400, "参数不正确")
						return
					}

					var row *User
					for _, item := range users {
						if item.ID == id {
							row = item
							break
						}
					}

					if row == nil {
						SendError(c, 404, "用户不存在")
						return
					}

					for i, item := range users {
						if item.ID == id {
							users = append(users[:i], users[i+1:]...)
							break
						}
					}

					stopBrook(row.Port)
					db.Delete(row)
					SendOk(c, nil)
				})
				r.NoRoute(func(c *gin.Context) {
					if strings.HasPrefix(c.Request.URL.Path, "/api") {
						director := func(req *http.Request) {
							req.URL.Scheme = "http"
							req.URL.Host = baseURL.Host
							req.Host = baseURL.Host
						}
						proxy := &httputil.ReverseProxy{Director: director}
						proxy.ServeHTTP(c.Writer, c.Request)
					} else if c.Request.Method == "GET" {
						c.File("public/index.html")
					}
				})
				logger.Infof("启动成功 http://%s:%d?token=%s", ip, config.Port, url.QueryEscape(config.Token))
				go func() {
					err := r.Run(fmt.Sprintf(":%d", config.Port))
					if err != nil {
						logger.Error("HTTP服务启动失败", err.Error())
					}
				}()
				go buildApp(ip, config.Port)
				cleanup()
				return nil
			},
		},
	}
	if err := app.Run(os.Args); err != nil {
		log.Fatal(err)
	}
}

func cleanup() {
	sigs := make(chan os.Signal, 1)
	signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM)
	<-sigs
	for _, v := range servers {
		_ = v.Shutdown()
	}
	if db != nil {
		db.Close()
	}
}

func Now() int64 {
	return time.Now().UnixNano() / 1e6
}

func reloadUsers(retry int) {
	// 用户
	data, err := httpGet("/node/userlist")
	if err != nil {
		if retry > 0 {
			time.Sleep(5 * time.Second)
			reloadUsers(retry - 1)
		}
		return
	}
	var rows = make([]*User, 0, 30)
	var rmap = make(map[int]bool)
	err = mapstructure.Decode(data, &rows)
	if err == nil {
		for _, user := range rows {
			_ = startBrook(user.Port, user.Passwd, user.Speed, user.Domain)
			rmap[user.Port] = true
		}
	}

	for _, user := range changes {
		if !rmap[user.Port] {
			stopBrook(user.Port)
		}
	}
	changes = rows
}

func startUser(user *User) {
	if user.Enable && user.U+user.D < user.Limit && user.ExpireAt > Now() {
		port2user_lock.Lock()
		port2user[user.Port] = user
		port2user_lock.Unlock()
		if debug {
			logger.Println(user.Port, user.Passwd, user.Speed, user.Domain)
		}
		err := startBrook(user.Port, user.Passwd, user.Speed, user.Domain)
		if err == nil && user.StartAt < 1 {
			user.StartAt = time.Now().UnixNano() / 1e6
		}
	} else if user.StartAt > 0 {
		user.StartAt = 0
		stopBrook(user.Port)
		port2user_lock.Lock()
		delete(port2user, user.Port)
		port2user_lock.Unlock()
	}
}

func stopBrook(port int) {
	lock.Lock()
	defer lock.Unlock()
	s := servers[port]
	if s != nil {
		_ = s.Shutdown()
		servers[port] = nil
	}
}

func startBrook(port int, passwd string, speed int, domain string) error {
	lock.Lock()
	defer lock.Unlock()
	brook.Traffic.LimitSpeed(port, speed)
	s := servers[port]
	if s != nil {
		var port0 int
		var passwd0 string
		var domain0 string
		switch v := s.(type) {
		case *brook.Server:
			port0 = v.TCPAddr.Port
			passwd0 = string(v.Password)
			domain0 = ""
		case *brook.WSServer:
			port0 = v.TCPAddr.Port
			passwd0 = string(v.Password)
			domain0 = "w" + v.Domain
		case *brook.SSServer:
			port0 = v.TCPAddr.Port
			passwd0 = string(v.Password)
			domain0 = "s"
		}
		if port0 == port && passwd0 == passwd && domain0 == domain {
			return nil
		}
		_ = s.Shutdown()
		servers[port] = nil
	}
	if err := CheckPort(port); err != nil {
		return err
	}
	var addr = fmt.Sprintf(":%d", port)
	var err error
	// 禁止brook模式，强制使用websocket模式
	if len(domain) < 1 {
		domain = "w"
	}
	if len(domain) > 0 && domain[0] == 'w' {
		// logger.Println("ws", addr, passwd, domain[1:])
		s, err = brook.NewWSServer(addr, passwd, domain[1:], 60, 0, 60)
		if err != nil {
			return err
		}
	} else if len(domain) > 0 && domain[0] == 's' {
		// logger.Println("ss", addr, passwd)
		s, err = brook.NewSSServer(addr, passwd, 60, 0, 60)
		if err != nil {
			return err
		}
	} else {
		// logger.Println("bk", addr, passwd)
		s, err = brook.NewServer(addr, passwd, 60, 0, 60)
		if err != nil {
			return err
		}
	}
	servers[port] = s
	go func() {
		err := s.ListenAndServe()
		if err != nil && !strings.Contains(err.Error(), "closed") {
			logger.Errorln("启动失败: ", err.Error())
			user := port2user[port]
			if user != nil {
				user.StartAt = 1
			}
		}
	}()
	return nil
}

func getIP(key string) string {
	bs, err := ioutil.ReadFile(".config.json")
	if err == nil {
		_ = json.Unmarshal(bs, config)
	}
	if key != "" {
		config.Key = key
	}
	if config.Key == "ZnJlZS1icm9va2lyZA==" {
		config.Key = "YnJvb2tpcmQ6Ly9zc3IuaW51MTI1NS5jbi9hcGk="
	}
	if config.Token == "" {
		config.Token = strconv.FormatInt(rand.Int63(), 36) + strconv.FormatInt(rand.Int63(), 36) + strconv.FormatInt(rand.Int63(), 36)
	}
	bs, err = json.MarshalIndent(config, "", "  ")
	if err == nil {
		err = ioutil.WriteFile(".config.json", bs, 0600)
		if err != nil {
			logger.Errorln("写入配置文件失败: ", err.Error())
		}
	}

	// 连接 master
	bs, err = base64.StdEncoding.DecodeString(config.Key)
	if err != nil {
		baseURL, err = url.Parse(strings.Trim(config.Key, "\n \t"))
	} else {
		baseURL, err = url.Parse(strings.Trim(string(bs), "\n \t"))
	}
	baseURL.RawQuery = "token=" + url.QueryEscape(baseURL.Scheme)
	baseURL.Scheme = "http"
	if err != nil {
		logger.Fatalln("错误的Key")
	}
	data, err := httpPost("/node/register", gin.H{
		"port": config.Port,
	})
	if err != nil {
		logger.Fatalln("错误的Key")
	}
	var info NodeRegisterData
	err = mapstructure.Decode(data, &info)
	if err != nil {
		logger.Errorln("无效的密钥")
	}
	return info.IP
}
